#!/usr/bin/env python

# Simulator for the prisoner switcharoo logic puzzle, as told be Cartalk

import sys
import random

class Actions:
  FlipA = 0
  FlipB = 1
  DeclareAllVisited = 2



class TomRay:  

  # rules for the Counter
  class Counter:
    _numOfPrisoners = 0
    _signalSet = False
    _count = 0

    def __init__(self, numOfPrisoners):
      self._numOfPrisoners = numOfPrisoners

    def DoSwitchRoom(self, switchPosA, switchPosB):
      if not switchPosA:
        return Actions.FlipB # throw away
      else:
        # switch moved, bump counter
        self._count += 1
  
        # check count, if sufficient make the declaration, otherwise reset sig
        if self._count == (self._numOfPrisoners - 1) * 2:
          return Actions.DeclareAllVisited
        else:
          return Actions.FlipA


  # rules for the Counted
  class Counted:
    _signalsSent = 0
    _seenSwitchBDown = False

    def DoSwitchRoom(self, switchPosA, switchPosB):
      if not switchPosA:
        if self._signalsSent < 2:
          self._signalsSent += 1
          return Actions.FlipA

      # throw away
      return Actions.FlipB


  def CreatePrisoner(self, prionerNumber, numOfPrisoners):
    if prionerNumber == 0:
      return TomRay.Counter(numOfPrisoners)
    else:
      return TomRay.Counted()

  def Name(self):
    return "Tom & Ray\'s slow method"


# She may not be the first to devise this solution, but for lack of a better name.
class Tabatha:

  _optimizedVersion = False

  def __init__(self, optimizedVersion):
    self._optimizedVersion = optimizedVersion


  # rules for the Counter
  class Counter:
    _numOfPrisoners = 0
    _signalSet = False
    _count = 0

    def __init__(self, numOfPrisoners):
      self._numOfPrisoners = numOfPrisoners

    def DoSwitchRoom(self, switchPosA, switchPosB):
      # bump count if the signal was set, and a is off
      if not switchPosA and self._signalSet:
          self._count += 1

      # signal set should be the opposite of the state of switch a,
      # since we will flip switch a
      self._signalSet = not switchPosA

      if self._count == self._numOfPrisoners - 1:
        return Actions.DeclareAllVisited
      else:
        return Actions.FlipA


  # rules for the Counted
  class Counted:
    _sentSignal = False
    _seenSwitchADown = False

    def DoSwitchRoom(self, switchPosA, switchPosB):
      if switchPosA:
        if self._seenSwitchADown and not self._sentSignal:
          self._sentSignal = True
          return Actions.FlipA
      else:
        self._seenSwitchADown = True

      # throw away
      return Actions.FlipB

  # rules for the optimized Counter
  # this Counter flips the dummy switch sometimes. There may be more cases he could flip
  # the dummy, but so far I've only let him flip it once everytime he doesnt see a signal on A
  class OptimizedCounter:
    _numOfPrisoners = 0
    _signalSet = False
    _count = 0

    _flippedLastTime = False

    def __init__(self, numOfPrisoners):
      self._numOfPrisoners = numOfPrisoners

    def DoSwitchRoom(self, switchPosA, switchPosB):
      # bump count if the signal was set, and a is off
      if not switchPosA and self._signalSet:
          self._count += 1

      if self._count == self._numOfPrisoners - 1:
        return Actions.DeclareAllVisited
      elif switchPosA:
        if self._flippedLastTime:
          self._flippedLastTime = False
          return Actions.FlipB # throw away this one
        else:
          self._flippedLastTime = True
          self._signalSet = False
          return Actions.FlipA
      else:
        self._flippedLastTime = True
        self._signalSet = True
        return Actions.FlipA # set the signal


  def CreatePrisoner(self, prionerNumber, numOfPrisoners):
    if prionerNumber == 0:
      if self._optimizedVersion:
        return Tabatha.OptimizedCounter(numOfPrisoners)
      else:
        return Tabatha.Counter(numOfPrisoners)
    else:
      return Tabatha.Counted()
  

  def Name(self):
    if self._optimizedVersion:
      return "My wife's faster method with additional optimization"
    else:
      return "My wife's faster method"



def PlayGame(strategy, numOfPrisoners, swAInitial, swBInitial):
  prisoners = []
  for i in range(0, numOfPrisoners):
    prisoners.append([strategy.CreatePrisoner(i, numOfPrisoners), False])

  swA = swAInitial
  swB = swBInitial

  visits = 0

  random.seed()  

  while True:
    visits += 1
    nextCon = random.randint(0, numOfPrisoners - 1)
    prisoners[nextCon][1] = True
    operation = prisoners[nextCon][0].DoSwitchRoom(swA, swB)
    if operation == Actions.FlipA:
      swA = not swA
    elif operation == Actions.FlipB:
      swB = not swB
    elif operation == Actions.DeclareAllVisited:
      break
    else:
      print 'Prisoner did not flip switch!'
      return -1

  # verify
  error = 0
  for i in range(0, numOfPrisoners):
    if not prisoners[i][1]:
      error += 1
      print 'prisoner ', i, ' did not go!'

  if error > 0:
    return -1

  return visits


def PlaySeveralGames(strategy, numOfPrisoners, numOfGames, swAInitial, swBInitial):
  totalVisits = 0
  for i in range(0, numOfGames):
    visits = PlayGame(strategy, numOfPrisoners, swAInitial, swBInitial)
    if visits < 0:
      print strategy.Name(), ' failed! All prisoners will be gator feed!'
      return -1

    totalVisits += visits

  print strategy.Name(), ' successful. Average visits needed for freedom: ', totalVisits / numOfGames
  return totalVisits


def PrintWinner(strat1, strat1Res, strat2, strat2Res):
  if strat1Res < strat2Res:
    winnerRes = strat1Res
    winnerStrat = strat1
    loserRes = strat2Res
    loserStrat = strat2
  else:
    winnerRes = strat2Res
    winnerStrat = strat2
    loserRes = strat1Res
    loserStrat = strat1

  print loserStrat.Name(), ' requires ', ((loserRes - winnerRes) / float(winnerRes)) * 100 ,'% more visits than ', winnerStrat.Name()
  print winnerStrat.Name(), ' required ', ((winnerRes) / float(loserRes)) * 100, '% of ', loserStrat.Name(), 'visits'

def PrintAvg(strat, res, runs):
  print strat.Name(), ' averaged ', float(res)/float(runs)

tr = TomRay()
tab = Tabatha(False)
tabOpt  = Tabatha(True)

prisoners = 23
runs = 1000

print 'Attempting ', runs, 'runs of each strategy, in each starting position.'

print 'Play with switch A initially off'
trRes = PlaySeveralGames(tr, prisoners, runs, False, False)
tabRes = PlaySeveralGames(tab, prisoners, runs, False, False)
tabOptRes = PlaySeveralGames(tabOpt, prisoners, runs, False, False)

if trRes < 0 or tabRes < 0 or tabOptRes < 0:
  print 'a strategy failed!'
  sys.exit(1)

print 'Play with switch A initially on'
trRes += PlaySeveralGames(tr, prisoners, runs, True, False)
tabRes += PlaySeveralGames(tab, prisoners, runs, True, False)
tabOptRes += PlaySeveralGames(tabOpt, prisoners, runs, True, False)

if trRes < 0 or tabRes < 0 or tabOptRes < 0:
  print 'a strategy failed!'
  sys.exit(1)

print 'All strategies successful'
print ''

PrintAvg(tr, trRes, runs * 2)
PrintAvg(tab, tabRes, runs * 2)
PrintAvg(tabOpt, tabOptRes, runs * 2)
print ''

PrintWinner(tab, tabRes, tr, trRes)
PrintWinner(tabOpt, tabOptRes, tr, trRes)
